import os, sys
import numpy as np
import cv2
import ncnn as pyncnn
from pathlib import Path
from glob import glob
import time

from utilities import *


class Inference:
    def __init__(self, model_path):
        self.model_path = model_path
        self.net = pyncnn.Net()
        w = Path(model_path)
        if not w.is_file():  # if not *.param
            w = next(w.glob("*.param"))  # get *.param file from *_ncnn_model dir
        self.net.load_param(str(w))
        self.net.load_model(str(w.with_suffix(".bin")))
        metadata = w.parent / "metadata.yaml"

        # Load external metadata YAML
        if isinstance(metadata, (str, Path)) and Path(metadata).exists():
            metadata = yaml_load(metadata)
        if metadata and isinstance(metadata, dict):
            for k, v in metadata.items():
                if k in {"stride", "batch"}:
                    metadata[k] = int(v)
                elif k in {"imgsz", "names", "kpt_shape"} and isinstance(v, str):
                    metadata[k] = eval(v)
            self.stride = metadata["stride"]
            self.task = metadata["task"]
            self.batch = metadata["batch"]
            self.imgsz = metadata["imgsz"]
            self.names = metadata["names"]
            self.kpt_shape = metadata.get("kpt_shape")

    def inference(self, origin):
        im0s = [origin]
        im = preprocess(im0s, self.imgsz, self.stride)
        b, ch, h, w = im.shape  # batch, channel, height, width
        # mat_in = pyncnn.Mat(im[0].cpu().numpy())

        mat_in = pyncnn.Mat(im[0])
        with self.net.create_extractor() as ex:
            ex.input(self.net.input_names()[0], mat_in)
            # WARNING: 'output_names' sorted as a temporary fix for https://github.com/pnnx/pnnx/issues/130
            y = []
            for x in sorted(self.net.output_names()):
                extract = ex.extract(x)
                y.append(np.array(extract[1])[None])
                # y = [np.array(ex.extract(x)[1])[None] for x in sorted(net.output_names())]

        preds = y[0]

        results = postprocess(preds, im, im0s, 0.25, 0.45)
        return results[0]


#############################################################
if __name__ == "__main__":
    # inf = Inference("weights/2in1_infrared_ncnn_model")
    inf = Inference(f"{os.getcwd()}/weights/yolo11n_ncnn_model")
    names = inf.names

    files = glob("images/*.jpg")
    for file in files:
        try:
            origin = cv2.imread(file)
            # origin = origin[160 : 160 + 320, 180 : 180 + 320]
        except Exception as e:
            print(f"Error reading image {file}: {e}")

        base_name = os.path.basename(file)

        if origin is None:
            continue

        # if origin is None:
        #     origin = cv2.imread(f"ncnn_inference/{file}")

        profiler = ops.Profile()
        with profiler:
            result = inf.inference(origin)

        print("run time:")
        print(profiler.dt)
        print("")

        im_show = origin
        for i, data in enumerate(result.data):
            cls, conf, bbox = result.cls[i], result.conf[i], result.data[i]
            # im_show = cv2.rectangle(origin, bbox, cls, conf, names)
            label = names[int(cls)]
            color = colors(int(cls), True)
            im_show = box_label(origin, bbox, label, conf, color)
        cv2.putText(im_show, f"{base_name}", (10, 30), cv2.FONT_HERSHEY_COMPLEX, 1, (50, 50, 220), 1)

        cv2.imshow("detect", im_show)
        cv2.waitKey(0)
